import tkinter as tk
from tkinter import ttk
from tkinter import filedialog
from PIL import Image, ImageTk
import cv2
import numpy as np
import mediapipe as mp
import threading
import time
import os
import io
import base64


class ImageOverlayApp:
    def __init__(self, root):
        self.root = root
        self.root.title("图片生成器")
        self.root.geometry("800x600")

        # 内置前景图 (base64编码的PNG)
        self.foreground_image = self.load_embedded_foreground()

        # 初始化变量
        self.background_path = ""
        self.output_path = ""

        # 创建GUI组件
        self.create_widgets()

        # 初始化MediaPipe
        self.mp_pose = mp.solutions.pose.Pose(
            static_image_mode=True, min_detection_confidence=0.5
        )

    def load_embedded_foreground(self):
        """加载内置的前景图"""
        # 这是一个示例透明PNG的base64编码（实际使用时替换为您的前景图）
        # return Image.open(io.BytesIO(base64.b64decode(base64_png.strip()))).convert("RGBA")
        return Image.open("foreground3.png").convert("RGBA")

    def create_widgets(self):
        # 顶部控制区域
        control_frame = tk.Frame(self.root)
        control_frame.pack(pady=10)

        # 背景图选择按钮
        self.bg_btn = tk.Button(
            control_frame, text="选择图片", command=self.select_background, width=15
        )
        self.bg_btn.pack(side=tk.LEFT, padx=5)

        # 处理按钮
        self.process_btn = tk.Button(
            control_frame,
            text="合成图片",
            command=self.start_processing,
            width=15,
            state=tk.DISABLED,
        )
        self.process_btn.pack(side=tk.LEFT, padx=5)

        # 保存按钮
        self.save_btn = tk.Button(
            control_frame,
            text="保存结果",
            command=self.save_result,
            width=15,
            state=tk.DISABLED,
        )
        self.save_btn.pack(side=tk.LEFT, padx=5)

        # 进度条
        self.progress = ttk.Progressbar(
            self.root, orient=tk.HORIZONTAL, length=400, mode="determinate"
        )
        self.progress.pack(pady=10)

        # 状态标签
        self.status_label = tk.Label(self.root, text="请选择图片", fg="blue")
        self.status_label.pack()

        # 图片显示区域
        self.image_frame = tk.Frame(self.root)
        self.image_frame.pack(pady=10)

        # 背景图预览
        self.bg_label = tk.Label(self.image_frame)
        self.bg_label.pack(side=tk.LEFT, padx=10)

        # 前景图预览
        fg_preview = self.foreground_image.copy()
        fg_preview.thumbnail((100, 100))
        fg_photo = ImageTk.PhotoImage(fg_preview)
        fg_label = tk.Label(self.image_frame, text="加个空调！")
        fg_label.pack(side=tk.LEFT)
        self.fg_preview_label = tk.Label(self.image_frame, image=fg_photo)
        self.fg_preview_label.image = fg_photo
        self.fg_preview_label.pack(side=tk.LEFT, padx=5)

        # 结果图预览
        self.result_label = tk.Label(self.image_frame)
        self.result_label.pack(side=tk.LEFT, padx=10)

    def select_background(self):
        file_path = filedialog.askopenfilename(
            title="选择图片", filetypes=[("图片文件", "*.jpg;*.jpeg;*.png")]
        )
        if file_path:
            self.background_path = file_path
            self.show_background_preview()
            self.process_btn.config(state=tk.NORMAL)
            self.status_label.config(text="背景图已加载", fg="green")

    def show_background_preview(self):
        try:
            img = Image.open(self.background_path)
            img.thumbnail((300, 300))
            photo = ImageTk.PhotoImage(img)
            self.bg_label.config(image=photo)
            self.bg_label.image = photo
        except Exception as e:
            self.status_label.config(text=f"预览错误: {str(e)}", fg="red")

    def show_result_preview(self, result_img):
        try:
            result_img.thumbnail((300, 300))
            photo = ImageTk.PhotoImage(result_img)
            self.result_label.config(image=photo)
            self.result_label.image = photo
        except Exception as e:
            self.status_label.config(text=f"结果预览错误: {str(e)}", fg="red")

    def start_processing(self):
        if not self.background_path:
            self.status_label.config(text="请先选择图片", fg="red")
            return

        # 禁用按钮
        self.process_btn.config(state=tk.DISABLED)
        self.bg_btn.config(state=tk.DISABLED)

        # 启动处理线程
        self.status_label.config(text="正在处理...", fg="blue")
        self.progress["value"] = 0
        threading.Thread(target=self.process_image, daemon=True).start()

    def process_image(self):
        try:
            # 模拟进度
            self.update_progress(10)

            # 加载背景图
            bg = Image.open(self.background_path).convert("RGBA")
            self.update_progress(20)

            # 使用MediaPipe检测人物
            bg_np = cv2.cvtColor(np.array(bg.convert("RGB")), cv2.COLOR_RGB2BGR)
            results = self.mp_pose.process(bg_np)
            self.update_progress(40)

            if not results.pose_landmarks:
                self.status_label.config(text="未检测到人物，使用默认位置", fg="orange")
                result = self.overlay_with_resize(bg, self.foreground_image)
            else:
                # 获取头部关键点坐标
                landmarks = results.pose_landmarks.landmark
                nose = landmarks[mp.solutions.pose.PoseLandmark.NOSE]
                left_ear = landmarks[mp.solutions.pose.PoseLandmark.LEFT_EAR]
                right_ear = landmarks[mp.solutions.pose.PoseLandmark.RIGHT_EAR]
                self.update_progress(50)

                # 计算头部中心位置
                head_center_x = int((nose.x + left_ear.x + right_ear.x) / 3 * bg.width)
                head_center_y = int(
                    (nose.y + min(left_ear.y, right_ear.y)) / 2 * bg.height
                )

                # 计算头部尺寸
                head_width = int(
                    max(abs(left_ear.x - right_ear.x) * bg.width * 1.8, bg.width // 5)
                )
                # head_height = int(head_width * 1.2)
                self.update_progress(60)

                # 调整前景图大小
                min_fg_size = int(head_width * 0.8)
                new_fg_width = max(
                    min_fg_size,
                    int(self.foreground_image.width * (head_width / bg.width)),
                )
                new_fg_height = int(
                    self.foreground_image.height
                    * (new_fg_width / self.foreground_image.width)
                )

                max_fg_size = bg.width // 2
                if new_fg_width > max_fg_size:
                    new_fg_width = max_fg_size
                    new_fg_height = int(
                        self.foreground_image.height
                        * (new_fg_width / self.foreground_image.width)
                    )

                fg_resized = self.foreground_image.resize(
                    (new_fg_width, new_fg_height), Image.LANCZOS
                )
                self.update_progress(70)

                # 计算位置
                position_x = head_center_x - new_fg_width // 2
                position_y = head_center_y - new_fg_height - int(new_fg_height * 0.6)

                # 边界检查
                position_x = max(0, min(position_x, bg.width - new_fg_width))
                position_y = max(0, min(position_y, bg.height - new_fg_height))
                self.update_progress(80)

                # 合并图片
                overlay = Image.new("RGBA", bg.size, (0, 0, 0, 0))
                overlay.paste(fg_resized, (position_x, position_y))
                result = Image.alpha_composite(bg, overlay)
                self.update_progress(90)

            self.result_image = result
            self.output_path = os.path.join(
                os.path.dirname(self.background_path),
                f"output_{os.path.basename(self.background_path)}",
            )

            self.update_progress(100)
            self.status_label.config(text="处理完成!", fg="green")
            self.show_result_preview(result)
            self.save_btn.config(state=tk.NORMAL)

        except Exception as e:
            self.status_label.config(text=f"处理错误: {str(e)}", fg="red")
        finally:
            self.process_btn.config(state=tk.NORMAL)
            self.bg_btn.config(state=tk.NORMAL)

    def overlay_with_resize(self, bg, fg):
        """未检测到人物时的默认处理"""
        new_width = bg.width // 3
        new_height = int(fg.height * (new_width / fg.width))
        fg_resized = fg.resize((new_width, new_height), Image.LANCZOS)

        position = ((bg.width - new_width) // 2, 0)

        overlay = Image.new("RGBA", bg.size, (0, 0, 0, 0))
        overlay.paste(fg_resized, position)
        return Image.alpha_composite(bg, overlay)

    def save_result(self):
        if not hasattr(self, "result_image"):
            self.status_label.config(text="没有可保存的结果", fg="red")
            return

        try:
            # 获取输出目录和文件名
            output_dir = os.path.dirname(self.background_path)
            base_name = os.path.basename(self.background_path)
            name, _ = os.path.splitext(base_name)

            # 默认保存为PNG格式（保留透明度）
            output_path = os.path.join(output_dir, f"output_{name}.png")

            # 检查是否需要透明通道
            if self.result_image.mode == "RGBA":
                self.result_image.save(output_path)
            else:
                self.result_image.convert("RGB").save(output_path)

            self.status_label.config(text=f"已保存到: {output_path}", fg="green")
            self.output_path = output_path  # 更新保存路径

        except Exception as e:
            self.status_label.config(text=f"保存失败: {str(e)}", fg="red")

    def update_progress(self, value):
        self.progress["value"] = value
        self.root.update_idletasks()
        time.sleep(0.1)  # 模拟处理延迟


if __name__ == "__main__":
    root = tk.Tk()
    app = ImageOverlayApp(root)
    root.mainloop()
